"""Ollama LLM Provider实现"""

import os
import time
from typing import Dict, Any, Optional, Generator, List
import numpy as np
import logging
import requests
from backend.llm.base import BaseLLMProvider, LLMProviderFactory

logger = logging.getLogger('llm.ollama')


class OllamaProvider(BaseLLMProvider):
    """Ollama LLM提供商实现"""
    
    def __init__(self, config: Dict[str, Any]):
        """初始化Ollama提供商
        
        Args:
            config: 配置参数，包含base_url、model等
        """
        super().__init__(config)
        self.base_url = config.get('base_url', os.getenv('OLLAMA_BASE_URL', 'http://localhost:11434'))
        self.model_name = config.get('model', os.getenv('OLLAMA_MODEL', 'llama3'))
        self.max_tokens = config.get('max_tokens', 4096)
        self.temperature = config.get('temperature', 0.7)
        self.top_p = config.get('top_p', 0.95)
        
        # 嵌入模型配置
        self.embedding_model_name = config.get('embedding_model', self.model_name)
        self.embedding_dimension = config.get('embedding_dimension', 4096)  # Ollama模型的嵌入维度通常较高
        
        # 请求超时设置
        self.request_timeout = config.get('timeout', 60)
    
    def load_model(self):
        """加载Ollama模型
        实际上是检查Ollama服务是否可用
        """
        try:
            # 检查Ollama服务是否可用 - 使用/api/tags端点替代不存在的/api/health
            tags_url = f"{self.base_url}/api/tags"
            response = requests.get(tags_url, timeout=self.request_timeout)
            
            if response.status_code == 200:
                try:
                    models_data = response.json()
                    available_models = [model['name'] for model in models_data.get('models', [])]
                    
                    # 验证配置的模型是否可用
                    if self.model_name in available_models:
                        self.model = self.model_name
                        logger.info(f"成功连接到Ollama服务，模型: {self.model_name} 可用")
                        return True
                    else:
                        logger.warning(f"Ollama服务可用，但配置的模型 {self.model_name} 不在可用模型列表中")
                        logger.warning(f"可用模型: {', '.join(available_models)}")
                        # 即使模型不在列表中，我们也尝试继续，因为某些情况下可能仍能工作
                        self.model = self.model_name
                        return True
                except Exception as json_error:
                    logger.warning(f"解析Ollama模型列表失败: {str(json_error)}")
                    # 即使解析失败，如果能连接到服务，我们也尝试继续
                    self.model = self.model_name
                    return True
            else:
                logger.error(f"Ollama服务检查失败，状态码: {response.status_code}")
                return False
        except requests.exceptions.ConnectionError:
            logger.error(f"无法连接到Ollama服务: {self.base_url}")
            self.model = None
            return False
        except Exception as e:
            logger.error(f"检查Ollama服务时发生异常: {str(e)}")
            self.model = None
            return False
    
    def generate(self, prompt: str, **kwargs) -> str:
        """生成文本响应
        
        Args:
            prompt: 输入提示
            **kwargs: 其他生成参数
        
        Returns:
            生成的文本
        """
        if not self.model:
            if not self.load_model():
                return "错误: Ollama服务未连接"
        
        try:
            # 准备参数
            params = {
                "model": self.model_name,
                "prompt": prompt,
                "stream": False,
                "options": {
                    "temperature": kwargs.get("temperature", self.temperature),
                    "top_p": kwargs.get("top_p", self.top_p),
                    "num_predict": kwargs.get("max_tokens", self.max_tokens),
                }
            }
            
            # 发送请求
            url = f"{self.base_url}/api/generate"
            response = requests.post(url, json=params, timeout=self.request_timeout)
            
            # 检查响应
            if response.status_code == 200:
                data = response.json()
                return data.get('response', '错误: 未获得有效响应')
            else:
                logger.error(f"Ollama生成请求失败，状态码: {response.status_code}, 响应: {response.text}")
                return f"错误: 请求失败，状态码: {response.status_code}"
        except Exception as e:
            logger.error(f"Ollama生成失败: {str(e)}")
            return f"错误: {str(e)}"
    
    def generate_stream(self, prompt: str, **kwargs) -> Generator[str, None, None]:
        """流式生成文本响应
        
        Args:
            prompt: 输入提示
            **kwargs: 其他生成参数
        
        Yields:
            文本片段
        """
        if not self.model:
            if not self.load_model():
                yield "错误: Ollama服务未连接"
                return
        
        try:
            # 准备参数
            params = {
                "model": self.model_name,
                "prompt": prompt,
                "stream": True,
                "options": {
                    "temperature": kwargs.get("temperature", self.temperature),
                    "top_p": kwargs.get("top_p", self.top_p),
                    "num_predict": kwargs.get("max_tokens", self.max_tokens),
                }
            }
            
            # 发送请求
            url = f"{self.base_url}/api/generate"
            with requests.post(url, json=params, stream=True, timeout=self.request_timeout) as response:
                
                # 检查响应状态
                if response.status_code != 200:
                    logger.error(f"Ollama流式生成请求失败，状态码: {response.status_code}")
                    yield f"错误: 请求失败，状态码: {response.status_code}"
                    return
                
                # 处理流式响应
                for line in response.iter_lines():
                    if line:
                        # 解析JSON行
                        try:
                            import json
                            chunk_data = json.loads(line.decode('utf-8'))
                            
                            # 提取内容
                            if 'response' in chunk_data and chunk_data['response']:
                                yield chunk_data['response']
                            
                            # 检查是否完成
                            if chunk_data.get('done', False):
                                break
                        except json.JSONDecodeError:
                            logger.warning(f"无法解析Ollama响应行: {line}")
        except Exception as e:
            logger.error(f"Ollama流式生成失败: {str(e)}")
            yield f"错误: {str(e)}"
    
    def generate_embedding(self, text: str) -> np.ndarray:
        """生成文本嵌入
        
        Args:
            text: 输入文本
        
        Returns:
            嵌入向量
        """
        if not self.model:
            if not self.load_model():
                return np.random.rand(self.embedding_dimension)
        
        try:
            # 准备参数
            params = {
                "model": self.embedding_model_name,
                "prompt": text
            }
            
            # 发送请求
            url = f"{self.base_url}/api/embeddings"
            response = requests.post(url, json=params, timeout=self.request_timeout)
            
            # 检查响应
            if response.status_code == 200:
                data = response.json()
                embedding = data.get('embedding', [])
                if embedding:
                    return np.array(embedding)
            else:
                logger.error(f"Ollama嵌入生成请求失败，状态码: {response.status_code}")
            
            # 如果出错，返回随机向量
            return np.random.rand(self.embedding_dimension)
        except Exception as e:
            logger.error(f"Ollama嵌入生成失败: {str(e)}")
            return np.random.rand(self.embedding_dimension)
    
    def generate_embeddings(self, texts: List[str]) -> List[np.ndarray]:
        """批量生成文本嵌入
        
        Args:
            texts: 输入文本列表
        
        Returns:
            嵌入向量列表
        """
        # Ollama API不支持批量嵌入，所以我们一个一个处理
        return [self.generate_embedding(text) for text in texts]
    
    def close(self):
        """关闭资源"""
        # Ollama是通过HTTP API调用的，不需要显式关闭
        self.model = None


# 注册Ollama提供商
LLMProviderFactory.register_provider('ollama', OllamaProvider)